import ast

from 程序核心代码.ast节点.ast操作器.AST操作器基类 import AST操作基类


class AST操作器(AST操作基类):
    _实例字典 = {}

    @classmethod
    def 获取实例(cls, ast树, *args, **kwargs):
        ast树_id = id(ast树)
        if ast树_id not in cls._实例字典:
            cls._实例字典[ast树_id] = cls(ast树, *args, **kwargs)
        return cls._实例字典[ast树_id]

    def __init__(self, ast树):
        if hasattr(self, '_已初始化'):
            return
        self._已初始化 = True
        super().__init__(ast树)

    def 获取关键词列表(self):
        return [node.id for node in ast.walk(self.ast树) if isinstance(node, ast.Name)]

    def 获取所有函数名(self):
        return [node.name for node in ast.walk(self.ast树) if isinstance(node, ast.FunctionDef)]

    def 获取所有变量名(self):
        return [node.id for node in ast.walk(self.ast树) if isinstance(node, ast.Name)]

    def 获取变量属性列表(self, 变量名 ,目标ast结构树=None):
        """获取指定变量的所有属性"""
        属性列表 = []
        目标ast结构树 = 目标ast结构树 or self.ast树
        for node in ast.walk(目标ast结构树):
            if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name):
                if node.value.id == 变量名:
                    属性列表.append(node.attr)
        return 属性列表

    def 查找节点(self, 条件):
        return [node for node in ast.walk(self.ast树) if 条件(node)]

    def 添加函数(self, 函数名, 参数列表, 函数体):
        新函数 = ast.FunctionDef(
            name=函数名,
            args=ast.arguments(
                posonlyargs=[],
                args=[ast.arg(arg=arg, annotation=None) for arg in 参数列表],
                kwonlyargs=[],
                kw_defaults=[],
                defaults=[]
            ),
            body=函数体,
            decorator_list=[]
        )
        self.ast树.body.append(新函数)
        self.记录操作(f"添加函数: {函数名}")

    def 删除函数(self, 函数名):
        self.ast树.body = [node for node in self.ast树.body if
                           not (isinstance(node, ast.FunctionDef) and node.name == 函数名)]
        self.记录操作(f"删除函数: {函数名}")

    def 捕获异常(self, 函数名):
        for node in self.ast树.body:
            if isinstance(node, ast.FunctionDef) and node.name == 函数名:
                try_stmt = ast.Try(
                    body=node.body,
                    handlers=[
                        ast.ExceptHandler(
                            type=ast.Name(id='Exception', ctx=ast.Load()),
                            name=ast.Name(id='e', ctx=ast.Store()),
                            body=[ast.Expr(ast.Call(func=ast.Name(id='print', ctx=ast.Load()),
                                                    args=[ast.Name(id='e', ctx=ast.Load())],
                                                    keywords=[]))]
                        )
                    ],
                    orelse=[],
                    finalbody=[]
                )
                node.body = [try_stmt]
                self.记录操作(f"在函数 {函数名} 中添加异常捕获")

    def 添加装饰器(self, 函数名, 装饰器):
        for node in self.ast树.body:
            if isinstance(node, ast.FunctionDef) and node.name == 函数名:
                decorator = ast.Name(id=装饰器, ctx=ast.Load())
                node.decorator_list.append(decorator)
                self.记录操作(f"为函数 {函数名} 添加装饰器: {装饰器}")

    def 替换变量名(self, 旧名, 新名):
        for node in ast.walk(self.ast树):
            if isinstance(node, ast.FunctionDef):
                for arg in node.args.args:
                    if arg.arg == 旧名:
                        arg.arg = 新名
            elif isinstance(node, ast.Name) and node.id == 旧名:
                node.id = 新名
        self.记录操作(f"替换变量名: {旧名} -> {新名}")

    def 添加类(self, 类名, 基类列表, 类体):
        新类 = ast.ClassDef(
            name=类名,
            bases=[ast.Name(id=基类, ctx=ast.Load()) for 基类 in 基类列表],
            keywords=[],
            body=类体,
            decorator_list=[]
        )
        self.ast树.body.append(新类)
        self.记录操作(f"添加类: {类名}")

    def 删除类(self, 类名):
        self.ast树.body = [node for node in self.ast树.body if
                           not (isinstance(node, ast.ClassDef) and node.name == 类名)]
        self.记录操作(f"删除类: {类名}")

    def 添加类方法(self, 类名, 方法名, 参数列表, 方法体):
        for node in self.ast树.body:
            if isinstance(node, ast.ClassDef) and node.name == 类名:
                新方法 = ast.FunctionDef(
                    name=方法名,
                    args=ast.arguments(
                        posonlyargs=[],
                        args=[ast.arg(arg='起始接口', annotation=None)] + [ast.arg(arg=arg, annotation=None) for arg in
                                                                       参数列表],
                        kwonlyargs=[],
                        kw_defaults=[],
                        defaults=[]
                    ),
                    body=方法体,
                    decorator_list=[]
                )
                node.body.append(新方法)
                self.记录操作(f"在类 {类名} 中添加方法: {方法名}")

    def 删除类方法(self, 类名, 方法名):
        for node in self.ast树.body:
            if isinstance(node, ast.ClassDef) and node.name == 类名:
                node.body = [method for method in node.body if
                             not (isinstance(method, ast.FunctionDef) and method.name == 方法名)]
                self.记录操作(f"在类 {类名} 中删除方法: {方法名}")

    def 添加属性(self, 类名, 属性名, 属性值):
        for node in self.ast树.body:
            if isinstance(node, ast.ClassDef) and node.name == 类名:
                属性节点 = ast.Assign(
                    targets=[ast.Name(id=属性名, ctx=ast.Store())],
                    value=ast.Constant(value=属性值)
                )
                node.body.insert(0, 属性节点)
                self.记录操作(f"在类 {类名} 中添加属性: {属性名} = {属性值}")


    def 删除属性(self, 类名, 属性名):
        for node in self.ast树.body:
            if isinstance(node, ast.ClassDef) and node.name == 类名:
                node.body = [attr for attr in node.body if
                             not (isinstance(attr, ast.Assign) and attr.targets[0].id == 属性名)]
                self.记录操作(f"在类 {类名} 中删除属性: {属性名}")

    def 查找类(self, 类名):
        return [node for node in self.ast树.body if isinstance(node, ast.ClassDef) and node.name == 类名]

    def 查找函数(self, 函数名):
        return [node for node in self.ast树.body if isinstance(node, ast.FunctionDef) and node.name == 函数名]

    def 获取操作日志(self):
        return self.操作日志

    def 获取变量路径(self, 变量名):
        路径列表 = []
        for node in ast.walk(self.ast树):
            if isinstance(node, ast.Name) and node.id == 变量名:
                路径列表.append(self._获取节点路径(self.ast树, node))
        return 路径列表

    def _获取节点路径(self, 父节点, 目标节点, 当前路径=None):
        if 当前路径 is None:
            当前路径 = []
        if isinstance(父节点, ast.AST):
            for field, value in ast.iter_fields(父节点):
                if value is 目标节点:
                    return 当前路径 + [field]
                elif isinstance(value, list):
                    for i, item in enumerate(value):
                        if item is 目标节点:
                            return 当前路径 + [field, i]
                        路径 = self._获取节点路径(item, 目标节点, 当前路径 + [field, i])
                        if 路径:
                            return 路径
                elif isinstance(value, ast.AST):
                    路径 = self._获取节点路径(value, 目标节点, 当前路径 + [field])
                    if 路径:
                        return 路径
        elif isinstance(父节点, list):
            for i, item in enumerate(父节点):
                if item is 目标节点:
                    return 当前路径 + [i]
                路径 = self._获取节点路径(item, 目标节点, 当前路径 + [i])
                if 路径:
                    return 路径
        return None
